
#include <armadillo>
class EP_tg_sep
{
	public:

	EP_tg_sep(void)
	{	
		boost::math::normal_distribution<> *d=new boost::math::normal_distribution<>(0,1); 
		_d=d;
	}
	~EP_tg_sep()
	{
		delete _d;
	}

	void operator()(mat m,mat Sig,mat b,int npass,double alpha=1)
	{
		int p=Sig.n_cols;
		mat mmi(p,1);
		mat Qmui(p,p);
		mat Q(p,p);
		Q=eye(p,p);
		cube Qi(p,p,p);
		mat Qmi(p,p);
		Qmui.fill(0);
		Qi.fill(0);
		mat e(p,1);
		mat L=chol(Sig);
		vec egp=eig_sym(Sig);
		vec eg=egp+1;
		//while(dist(eg,egp)>0.01)
		for(int k=0;k<npass;k++)
		{
			egp=eg;
			for(int i=0;i<p;i++)
			{
				mat oldmu=m;
				mat oldQ=Q;
				mat temp=Qi(span(i,i),span::all,span::all);
				Qmi=Q-temp;
				mmi=inv(Qmi)*(Q*m-Qmui(i,span::all).t());
				e.fill(0);
				e(i,0)=1;
				mat Vmi=inv(Qmi);
				cout << "Vmi " <<Vmi;
				mat G=L(i,span::all);
				double norm2=sqrt(as_scalar(G*Vmi*G.t()));
				mat dlZm=(1/(1-cdf(*_d,as_scalar((b(i,0)-dot(G,mmi))/norm2)))*pdf(*_d,as_scalar((b(i,0)-dot(G,mmi))/norm2)))*G/norm2;	
				mat p2=0.5*(G.t()*G)*(b(i,0)-dot(G,mmi))/pow(norm2,3);
				mat dlZS=(1/(1-cdf(*_d,as_scalar((b(i,0)-dot(G,mmi))/norm2)))*pdf(*_d,as_scalar((b(i,0)-dot(G,mmi))/norm2)))*p2;
				m=mmi+Vmi*dlZm.t();
				cout << "m " << m;
				mat V=Vmi-Vmi*(dlZm.t()*dlZm-2*dlZS)*Vmi;
				cout << "V " << V;
				Q=inv(V);
				mat T=Qi(span(i,i),span::all,span::all);
				Qi(span(i,i),span::all,span::all)=alpha*(Q-Qmi)+(1-alpha)*T;;
			//	temp=Qi(span(i,i),span::all,span::all);
				Qmui(span::all,i)=(alpha*(Q*m-Qmi*mmi))-(1-alpha)*Qmui(span::all,i);
			}
				egp=eig_sym(inv(Q));
		
		}
		_m=m;
		_S=inv(Q);
		//cout << _S;

	}
	double dist(vec x1, vec x2){
		vec x3=x1-x2;
		return dot(x3,x3);
	}
	mat Get_Sig(void){return _S;}
	mat Get_Mu(void){return _m;}
	private:
	mat _m;
	mat _S;
	boost::math::normal_distribution<> *_d; 
};

